open Filename
open Printf
open Next
open Circuit
open Str


let qfile = ref ""
let evfile = ref ""
let network = ref ""
let schemafile = ref ""
let modelfile = ref ""
let outfile = ref ""
let usage = "Usage: spquery -i <query> -m <model_directory> [...]"
let args = Arg.align
  ([
    ("-q",  Arg.Set_string qfile,  " Query file");
    (*("-net", Arg.Set_string network, " Dataset Network");*)
    (*("-s",  Arg.Set_string schemafile,  " Schema file"); *)
    ("-ev",  Arg.Set_string evfile,  " Evidence file");
    ("-o",  Arg.Set_string outfile,  " Inference result file");
    ("-m",  Arg.Set_string modelfile,  " Input SPN (SPAC) model directory")] 
    @ Ext.common_arguments)



exception HSplitNotPossible
exception VSplitNotPossible
exception NodeTypeError


type ndetails = TimesNode
        | PlusNode
        | NullNode 
  | ACNode;;


type spnode = {
  mutable id: int;
  mutable parents: int array;
  mutable children: int array;
  mutable params: float array;
  mutable nodetype: ndetails;
  mutable schema: (int * int) array;
  mutable data: int array array;
  mutable acname: string;
  mutable ac_index: int ;
  mutable counts: int array;
  mutable activated_link: int;
  mutable kl: float;
  mutable thresh: float;
}

(*let root = ref spnode
*)


let node_array = ref [||]
let comps = ref [||] 
let logpTemp = ref [||]
let logeTemp = ref [||]
let feature_lists = ref [||]
let scratches = ref [||] 
let logzs = ref [||] 
(*let data_parts = ref [||]
*)
let get_next_id () = 
  Array.length !node_array



let create_node t pt c pm d s=
  let n = {id = get_next_id(); 
  parents = pt;
  children = c; 
  params = pm;
  nodetype = t;
  schema = s; 
  data = d;
  acname = "";
  ac_index = 0;
  counts = [||];
  activated_link = 0;
  kl = 0.0;
  thresh = 0.0} in
  node_array := Array.append !node_array [|n|];
  n


let create_acnode = create_node ACNode [||] [||] [||]  

let create_nullnode = create_node NullNode [||] [||] [||]  
 
(*
let append_node node =

  ignore()
*)

let create_plus = create_node PlusNode [||] [||] [||] 


let create_times  = create_node TimesNode [||] [||] [||] 


let len a = Array.length a 

let node_as_array n = [|n|]

let node_as_list n = [n]

let add_child node child = 
  node.children <- Array.append node.children  [|child.id|];
  node.params <- Array.append node.params [|log(float_of_int(Array.length child.data)) -. log( float_of_int(Array.length node.data))|];
  node.counts <- Array.append node.counts [|(Array.length child.data)|];
  ignore()

let set_parent node parent =
  add_child parent node;
  node.parents <- Array.append node.parents [|parent.id|];
  ignore()

let print_params node =
  Array.iter (fun p->printf "%f " (exp p)) node.params 

let rec print_tree node depth nt = 
  for i=0 to depth do
    printf "  "
  done;
  (*printf "id: %d sample: %d vars: %d type: " node.id (Array.length node.data) (Array.length node.schema);*)
  printf "id: %d vars: %d type: " node.id (*(Array.length node.data)*) (Array.length node.schema);
  (match node.nodetype with ACNode-> printf "ac\n" | TimesNode-> printf "* \n" (* "kl: %f, thr: %f\n" (node.kl/.nt) node.thresh*) | PlusNode-> printf "+ "; print_params node; printf "\n" (*; printf "kl: %f\n" (node.kl/.nt)*) | NullNode -> ignore());
  let n = len node.children in
  for i=0 to n - 1 do
    print_tree !node_array.(node.children.(i)) (depth+1) nt
  done;
  ignore()


let print_spn node nt =
  print_tree node 0 nt



exception CircuitLoadException

let get_node_type t = match t with ACNode->"ac" | PlusNode-> "+"| TimesNode ->"*" | NullNode->"null"

let load_circuit filename =
  try Circuit.load_with_features (open_in filename) 
  with Circuit.Unknown_child(n,c) -> printf "Error: node %s references unknown child node %s.\n" n c; raise CircuitLoadException

let load_comps () =
  printf "load comps..%d\n"(Array.length !node_array) ; flush_all();
  feature_lists := [||];
  comps := [||];
  for i=1 to (Array.length !node_array) - 1  do
  (*  printf "node id:%d type:%s\n" !node_array.(i).id (match !node_array.(i).nodetype with ACNode->"ac" | PlusNode-> "+"| TimesNode ->"*");
  *)
    let nodetype = !node_array.(i).nodetype in
    if (nodetype == ACNode) then (* || nodetype == PlusNode || nodetype == TimesNode) then *)
    begin
      let comp_name = !node_array.(i).acname in
    (*  printf "%s\n" comp_name; flush_all();*)
	  let full_comp_name = String.concat dir_sep [!modelfile; comp_name] in
      let (component, fl) = load_circuit full_comp_name in
      let n = Array.length !comps in 
      comps := Array.append !comps [|component|];
      feature_lists := Array.append !feature_lists [|fl|];
      !node_array.(i).ac_index <- n;
    end
  done;
  ignore()



let answer_acnode_cached node_index i  =
  let ac_index = !node_array.(node_index).ac_index in
  !logpTemp.(ac_index).(i)


let answer_acnode_cached_ev node_index i  =
  let ac_index = !node_array.(node_index).ac_index in
  let p = !logeTemp.(ac_index).(i) in
  p

let answer_acnode_query node_index q = 
  let ac_index = !node_array.(node_index).ac_index in
  let circ = !comps.(ac_index) in
  let scratch = !scratches.(ac_index) in 
  let logz = !logzs.(ac_index) in
  let schema_indexes = Array.map (snd) (!node_array.(node_index).schema ) in
  let c_q = Array.take q schema_indexes in
  let logp = Circuit.logprob_x scratch circ c_q in
  let final = logp -. logz in
  final


let get_type_name nodetype = 
  if nodetype == PlusNode then "+" else if nodetype == TimesNode then "*" else "ac"

let normalize a =
  let sum = Array.sumf (Array.map (exp) a) in 
  let n_a = Array.map (fun x->log ((exp x) /. sum) ) a in
  n_a



let rec answer_spac_qurey node q =  
  let children = node.children in
  let n = Array.length children in 
  if node.nodetype == PlusNode then
  begin
    (*let n_params = Array.normalizelog node.params in*)
    let n = Array.length children in 
    let logprobs = Array.make n 0.0 in
    for i=0 to n - 1 do
      logprobs.(i) <- answer_spac_qurey !node_array.(children.(i)) q
    done;
    let final_probs = Array.map2 (fun logp logw-> logp +. logw ) logprobs node.params in 

    (*let total = float_of_int(Array.sum node.counts) in*)
    let final_result = (Ext.alogsumexp final_probs) in 
    final_result
  end
  else if node.nodetype == TimesNode then
  begin
    let logprobs = Array.make n 0.0 in
    for i=0 to (Array.length children) - 1 do
      logprobs.(i) <- answer_spac_qurey !node_array.(children.(i)) q
    done; 
    (Array.sumf logprobs) 
  end 
  else if node.nodetype == ACNode then
    answer_acnode_query node.id q 

  else raise NodeTypeError  

let rec answer_spac_cached_ev node q_i =
  (*printf "node.id %d\n" node.id;*)
  let children = node.children in
  let n = Array.length children in 
  if node.nodetype == PlusNode then
  begin
    let n = Array.length children in 
    let logprobs = Array.make n 0.0 in
    for i=0 to n - 1 do
      logprobs.(i) <- answer_spac_cached_ev !node_array.(children.(i)) q_i
    done;
    let final_probs = Array.map2 (fun logp logw-> logp +. logw ) logprobs node.params in 
    (*let total = float_of_int(Array.sum node.counts) in*)
    let final_result = (Ext.alogsumexp final_probs) in
    final_result
  end
  else if node.nodetype == TimesNode then
  begin
    let logprobs = Array.make n 0.0 in
    for i=0 to (Array.length children) - 1 do
      logprobs.(i) <- answer_spac_cached_ev !node_array.(children.(i)) q_i
    done; 
    let final_result = (Array.sumf logprobs) in
    final_result
  end 
  else if node.nodetype == ACNode then
    answer_acnode_cached_ev node.id q_i

  else raise NodeTypeError  


let rec answer_spac_cached node q_i = 
  let children = node.children in
  let n = Array.length children in 
  (*printf "id: %d answer query: %d\n" node.id q_i;
  flush_all();*)
  if node.nodetype == PlusNode then
  begin
    (*let n_params = Array.normalizelog node.params in*)
    let n = Array.length children in 
    let logprobs = Array.make n 0.0 in
    for i=0 to n - 1 do
      logprobs.(i) <- answer_spac_cached !node_array.(children.(i)) q_i
    done;
    let final_probs = Array.map2 (fun logp logw-> logp +. logw ) logprobs node.params in 
    (*let total = float_of_int(Array.sum node.counts) in*)
    let final_result = (Ext.alogsumexp final_probs) in

    (*printf "node: %d p = %f\n" node.id (exp final_result);
	*)
	final_result
  end
  else if node.nodetype == TimesNode then
  begin
    let logprobs = Array.make n 0.0 in
    for i=0 to (Array.length children) - 1 do
      logprobs.(i) <- answer_spac_cached !node_array.(children.(i)) q_i
    done; 
    let final_result = (Array.sumf logprobs) in
       (*printf "node: %d p = %f\n" node.id (exp final_result);
       *)
	   final_result
  end 
  else if node.nodetype == ACNode then begin
	let final_result = answer_acnode_cached node.id q_i in
  	
    (*printf "node: %d q_i: %d p = %f\n" node.id q_i (exp final_result);
	*)
	final_result
  end
  else raise NodeTypeError  




let answer_query root q =
  let logp = answer_spac_qurey root q in
  logp


let answer_query_cached root q_i =
  let logp = answer_spac_cached root q_i in
  logp

let answer_query_cached_ev root q_i =
  let logp = answer_spac_cached_ev root q_i in
  logp

exception InsufficientArguments


let print_node out node =
  let n = (Array.length node.children) in
  fprintf out "%d %s\n" node.id  (match node.nodetype with ACNode-> "ac" | TimesNode-> "*" | PlusNode-> "+" | NullNode ->  "null");
  for i=0 to n-1 do
    fprintf out "%d " node.children.(i)
  done;
  fprintf out "\n";
  for i=0 to n-1 do
    fprintf out "%f " node.params.(i)
  done;
  fprintf out "\n"

let print_model out root =
  fprintf out "%d\n" (Array.length !node_array);
  Array.iter (print_node out) !node_array

exception ParseError of string
let space = regexp " " 

let read_line_tokens model = split space (input_line model) 

let answer_joint_queries circ scratch logz queries evidence =
  let n = Array.length queries in
 
  let logps = Array.make n 0.0 in
  for k=0 to n - 1 do
      (* By definition of conditional probability:
           log P(q|e) = log P(q) - log P(e) *)
      let q = queries.(k) in

    (* Input validation *)
      Data.check_point circ.Circuit.schema q;

      let e = evidence.(k) in
      (* Check for compatibility between evidence and query. *)
      if Array.length e > 0 then
        Data.check_point circ.Circuit.schema e;
      let mismatch = ref false in
      for i = 0 to Array.length e - 1 do
        if e.(i) >= 0 then begin
          if q.(i) < 0 then 
            q.(i) <- e.(i)
          else if q.(i) != e.(i) then
            mismatch := true
        end
      done;

      (* Compute probability of query and evidence *)
      let logp_q = if !mismatch then log 0. 
                   else Circuit.logprob_x scratch circ q in
      let logp_ev = Circuit.logprob_x scratch circ e in

      let l = logp_q -. logp_ev in
      logps.(k) <- l 
  done;
  logps

let parse_node model = 
  let tokens = split space (input_line model) in
  (*printf "line: %s %s %s\n" (List.nth tokens 0) (List.nth tokens 1) (List.nth tokens 2); 
  *)
  let i = match (List.nth tokens 0) with "n" -> 0 | _-> raise (ParseError "not a node") in
  ignore(i);
  let nodeid = int_of_string (List.nth tokens 1) in
  let nodetype = match (List.nth tokens 2) with "+" -> PlusNode | "ac" -> ACNode | "*" -> TimesNode | "null"->NullNode | _-> raise(ParseError "node type incorrect") in
  let parentid = int_of_string (List.nth tokens 3) in
  (*printf "network: %s nodeid :%d nodetype:%s %s parentid :%d \n" !network nodeid (List.nth tokens 2) (get_node_type(nodetype)) parentid;
  *)
  let parents = if parentid < 0 then [||] else [|parentid|] in
  (match nodetype with PlusNode -> 
    begin
    let tokens = read_line_tokens model in
    let children = Array.of_list (List.map (int_of_string) tokens) in
    
    let tokens = read_line_tokens model in
    let params = Array.of_list (List.map (float_of_string) tokens) in
    let tokens = read_line_tokens model in
    let schema = Array.of_list (List.map (int_of_string) tokens ) in
    let i_schema = Array.map (fun s->(2,s))schema in
    let n = create_node  nodetype parents children params [||] i_schema in
    (*!node_array.(nodeid).acname <- sprintf "%s-%d.ac" !network nodeid;*)
    ignore(n)
    end
    | TimesNode ->
    begin

    let tokens = read_line_tokens model in
    let thresh = float_of_string(List.hd tokens) in
    let tokens = read_line_tokens model in
    let children = Array.of_list (List.map (int_of_string) tokens) in

    let tokens = read_line_tokens model in
    let schema = Array.of_list (List.map (int_of_string) tokens ) in
    let i_schema = Array.map (fun s->(2,s))schema in
    let n = create_node  nodetype parents children [||] [||] i_schema in
    n.thresh<-thresh;
    ignore(n)
    (*!node_array.(nodeid).acname <- sprintf "%s-%d.ac" !network nodeid;*)
    end
    | ACNode ->
    begin


    let tokens = read_line_tokens model in
    let schema = Array.of_list (List.map (int_of_string) tokens ) in
    let i_schema = Array.map (fun s->(2,s))schema in
    let ac_node = create_node nodetype parents [||] [||] [||] i_schema in
   
   	ignore(ac_node);
    !node_array.(nodeid).acname <- sprintf "%s-%d.ac" !network nodeid;

    end
    | NullNode-> begin let nullnode=create_nullnode [||] [||] in ignore(nullnode) end
   );
  
  ignore()
  (*!node_array.(nodeid).nodetype <- nodetype;*)
  (*!node_array.(nodeid).acname <- sprintf "%s-%d.ac" !network nodeid*)

let load_model modelfile =
  try
    let model = open_in modelfile in
    network := input_line model;
    let num_nodes = int_of_string (input_line model) in
    printf "Num of nodes: %d\n" num_nodes;  
    for i=1 to num_nodes do
      parse_node model;     
    done;
  with Failure _ -> printf "ERROR\n" 

let preload_query_results queries = 
  let n = Array.length queries in
  logpTemp := Array.make (Array.length !comps) [||];  
  for i = 0 to (Array.length !node_array) - 1 do
    let node = !node_array.(i) in 
    if node.nodetype == ACNode (* || node.nodetype == PlusNode || node.nodetype == TimesNode*) then 
    begin
      let result = Array.make n 0.0 in
      let indeces = Array.map (snd) node.schema in
      let myqueries = Array.takeCol queries indeces in
  
      let ac_index = node.ac_index in
      let circ = !comps.(ac_index) in
      let scratch = !scratches.(ac_index) in 
      let logz = !logzs.(ac_index) in

      for k=0 to n - 1 do
        let logp = Circuit.logprob_x scratch circ myqueries.(k) in
        result.(k) <- (logp -. logz) 
      done;
      !logpTemp.(ac_index) <- result
    end
  done


let preload_ev_results queries = 
  let n = Array.length queries in
  logeTemp := Array.make (Array.length !comps) [||];  
  for i = 0 to (Array.length !node_array) - 1 do
    let node = !node_array.(i) in 
    if node.nodetype == ACNode (* || node.nodetype == PlusNode || node.nodetype == TimesNode*) then 
    begin
      let result = Array.make n 0.0 in
      let indeces = Array.map (snd) node.schema in
      let myqueries = Array.takeCol queries indeces in
  
      let ac_index = node.ac_index in
      let circ = !comps.(ac_index) in
      let scratch = !scratches.(ac_index) in 
      let logz = !logzs.(ac_index) in

      for k=0 to n - 1 do
        let logp = Circuit.logprob_x scratch circ myqueries.(k) in
        result.(k) <- (logp -. logz)  
      done;
      !logeTemp.(ac_index) <- result
    end
  done

let preload_with_ev queries evidence = 
  logpTemp := Array.make (Array.length !comps) [||];  
  for i = 0 to (Array.length !node_array) - 1 do
    let node = !node_array.(i) in 
    if node.nodetype == ACNode (* || node.nodetype == PlusNode || node.nodetype == TimesNode*) then 
    begin
      let indeces = Array.map (snd) node.schema in
      let myqueries = Array.takeCol queries indeces in
      let myevidence = Array.takeCol evidence indeces in
  
      let ac_index = node.ac_index in
      let circ = !comps.(ac_index) in
      let scratch = !scratches.(ac_index) in 
      let logz = !logzs.(ac_index) in
  
      let result = answer_joint_queries circ scratch logz  myqueries myevidence in
      printf "acname: %s \n" node.acname;
      Array.iter (printf "%f\n") result;
      !logpTemp.(ac_index) <- result
    end
  done
(* main main *) 


let endswith s1 s2 =
  let re = Str.regexp (Str.quote s2 ^ "$") in
  try ignore (Str.search_forward re s1 0); true
  with Not_found -> false

let _ =
  Arg.parse args ignore usage;
  if  !qfile = "" || !modelfile = "" then (Arg.usage args usage; exit (-1));
  Ext.common_log_init ();
  
  if endswith !modelfile dir_sep then
    modelfile := String.sub !modelfile 0 ((String.length !modelfile) - 1);

  if not (check_suffix !modelfile ".spn") then begin printf "Only accepts .spn as the input model file"; exit 1 end;
	

  let test = Data.load_data_ar !qfile in
  (*let schema = Data.load_schema !schemafile in*)
  let t1 = Sys.time() in
  ignore(t1);
  printf "loading the model .... \n";
  load_model (String.concat dir_sep [!modelfile; "spac.m"]);  
  
  load_comps();
  
  
  scratches := Array.map (Circuit.create_scratch) !comps;
  logzs := Array.map2 (fun s c->Circuit.compute_z s c ) !scratches !comps;
  
  let t3 = Sys.time() in
  printf "preload test data into ACs..\n";
  flush_all();
  
  let root = !node_array.(0) in 
  let finallogprobs = ref [||] in
  if !evfile <> "" then begin
    let ev = Data.load_data_ar !evfile in
    (*preload_with_ev test ev; *)

	preload_ev_results ev;
    preload_query_results test;
    
    printf "Answering test queries..\n";
    
	let testlogprobs = Array.mapi (fun i x-> answer_query_cached root i) test in
    let evlogprobs = Array.mapi (fun i x-> answer_query_cached_ev root i) ev in
    finallogprobs := Array.map2 (fun q e -> q-.e) testlogprobs evlogprobs 
  
  end else begin
    preload_query_results test;
    printf "Answering test queries..\n";
    finallogprobs := Array.mapi (fun i x-> answer_query_cached root i) test
  end;

  let llg_test = Array.sumf !finallogprobs in
  let nt = float_of_int ( Array.length test) in
  let avg_test = llg_test /. nt in
  let t4 = Sys.time() in
  ignore(t4);ignore(t3);

  print_spn root nt;

  printf "result avg llg: %.7f\n" avg_test; 

	
  if !outfile <> "" then
  begin
    let inference_channel = open_out !outfile in
    Array.iter (fprintf inference_channel "%.7f\n" ) !finallogprobs;
    fprintf inference_channel "avg = %f\n" avg_test;
	close_out_noerr inference_channel
  end;
  (*
  if !outputmodel <> "" then 
  begin
    let model_channel = open_out !outputmodel in
    print_model model_channel root;
    close_out_noerr model_channel
  end;


  printf "result %s %d %f %f %f %f %f %f %d %f %f %f %f %f\n" !network !num_part !split_penalty !l1param !stddev !ratio !cut_thresh !lambda !maxe avg_data avg_valid avg_test (t2 -. t1) (t4 -. t3);
  ignore()

  *)

